/*

 * Licensed to the Apache Software Foundation (ASF) under one

 * or more contributor license agreements.  See the NOTICE file

 * distributed with this work for additional information

 * regarding copyright ownership.  The ASF licenses this file

 * to you under the Apache License, Version 2.0 (the

 * "License"); you may not use this file except in compliance

 * with the License.  You may obtain a copy of the License at

 *

 *     http://www.apache.org/licenses/LICENSE-2.0

 *

 * Unless required by applicable law or agreed to in writing, software

 * distributed under the License is distributed on an "AS IS" BASIS,

 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.

 * See the License for the specific language governing permissions and

 * limitations under the License.

 */

package com.bff.gaia.unified.sdk.loadtests;



import com.amazonaws.regions.Regions;

import com.bff.gaia.unified.sdk.Pipeline;

import com.bff.gaia.unified.sdk.coders.ByteArrayCoder;

import com.bff.gaia.unified.sdk.coders.CoderException;

import com.bff.gaia.unified.sdk.coders.KvCoder;

import com.bff.gaia.unified.sdk.io.Read;

import com.bff.gaia.unified.sdk.io.gcp.pubsub.PubsubIO;

import com.bff.gaia.unified.sdk.io.gcp.pubsub.PubsubMessage;

import com.bff.gaia.unified.sdk.io.kafka.KafkaIO;

import com.bff.gaia.unified.sdk.io.kinesis.KinesisIO;

import com.bff.gaia.unified.sdk.io.synthetic.SyntheticBoundedSource;

import com.bff.gaia.unified.sdk.io.synthetic.SyntheticOptions;

import com.bff.gaia.unified.sdk.io.synthetic.SyntheticSourceOptions;

import com.bff.gaia.unified.sdk.options.*;

import com.bff.gaia.unified.sdk.transforms.MapElements;

import com.bff.gaia.unified.sdk.transforms.SimpleFunction;

import com.bff.gaia.unified.sdk.values.KV;

import com.bff.gaia.unified.sdk.values.PCollection;

import com.bff.gaia.unified.sdk.options.ApplicationNameOptions;

import com.bff.gaia.unified.sdk.options.Description;

import com.bff.gaia.unified.sdk.options.PipelineOptions;

import com.bff.gaia.unified.sdk.options.PipelineOptionsFactory;

import com.bff.gaia.unified.sdk.options.Validation;

import com.bff.gaia.unified.sdk.util.CoderUtils;

import org.apache.kafka.common.serialization.StringSerializer;



import java.io.IOException;

import java.util.Arrays;

import java.util.HashMap;

import java.util.Map;



import static java.nio.charset.StandardCharsets.UTF_8;

import static com.bff.gaia.unified.sdk.util.CoderUtils.encodeToByteArray;



/**

 * Pipeline that generates synthetic data and publishes it in a PubSub or Kafka topic or in a

 * Kinesis stream.

 *

 * <p>To run it manually, use the following command:

 *

 * <pre>

 *  ./gradlew :unified-sdks-java-load-tests:run -PloadTest.args='

 *    --pubSubTopic=TOPIC_NAME

 *    --kafkaBootstrapServerAddress=SERVER_ADDRESS

 *    --kafkaTopic=KAFKA_TOPIC_NAME

 *    --sourceOptions={"numRecords":1000,...}

 *    --kinesisStreamName=KINESIS_STREAM_NAME

 *    --kinesisPartitionKey=KINESIS_PARTITION_KEY

 *    --kinesisAwsKey=KINESIS_AWS_KEY

 *    --kinesisAwsSecret=KINESIS_AWS_SECRET

 *    --kinesisAwsRegion=KINESIS_AWS_REGION'

 *    -PloadTest.mainClass="com.bff.gaia.unified.sdk.loadtests.SyntheticDataPublisher"

 *  </pre>

 *

 * <p>If parameters related to a specific sink are provided (Kafka, PubSub or Kinesis), the pipeline

 * writes to the sink. Writing to more than one sink is also acceptable.

 */

public class SyntheticDataPublisher {



  private static final KvCoder<byte[], byte[]> RECORD_CODER =

      KvCoder.of(ByteArrayCoder.of(), ByteArrayCoder.of());



  private static Options options;



  /** Options for the pipeline. */

  public interface Options extends PipelineOptions, ApplicationNameOptions {



    @Description("Options for synthetic source")

    @Validation.Required

    String getSourceOptions();



    void setSourceOptions(String sourceOptions);



    @Description("PubSub topic to publish to")

    String getPubSubTopic();



    void setPubSubTopic(String topic);



    @Description("Kafka server address")

    String getKafkaBootstrapServerAddress();



    void setKafkaBootstrapServerAddress(String address);



    @Description("Kafka topic")

    String getKafkaTopic();



    void setKafkaTopic(String topic);



    @Description("Kinesis partition key")

    String getKinesisPartitionKey();



    void setKinesisPartitionKey(String partitionKey);



    @Description("Kinesis partition name")

    String getKinesisStreamName();



    void setKinesisStreamName(String kinesisStreamName);



    @Description("AWS secret for Kinesis")

    String getKinesisAwsSecret();



    void setKinesisAwsSecret(String kinesisAwsSecret);



    @Description("AWS key for Kinesis")

    String getKinesisAwsKey();



    void setKinesisAwsKey(String kinesisAwsKey);



    @Description("AWS region for Kinesis")

    String getKinesisAwsRegion();



    void setKinesisAwsRegion(String kinesisAwsRegion);

  }



  public static void main(String[] args) throws IOException {

    options = PipelineOptionsFactory.fromArgs(args).withValidation().as(Options.class);



    SyntheticSourceOptions sourceOptions =

        SyntheticOptions.fromJsonString(options.getSourceOptions(), SyntheticSourceOptions.class);



    Pipeline pipeline = Pipeline.create(options);

    PCollection<KV<byte[], byte[]>> syntheticData =

        pipeline.apply("Read synthetic data", Read.from(new SyntheticBoundedSource(sourceOptions)));



    if (options.getKafkaBootstrapServerAddress() != null && options.getKafkaTopic() != null) {

      writeToKafka(syntheticData);

    }

    if (options.getPubSubTopic() != null) {

      writeToPubSub(syntheticData);

    }

    if (allKinesisOptionsConfigured()) {

      writeToKinesis(syntheticData);

    }

    pipeline.run().waitUntilFinish();

  }



  private static boolean allKinesisOptionsConfigured() {

    return options.getKinesisAwsKey() != null

        && options.getKinesisAwsRegion() != null

        && options.getKinesisAwsSecret() != null

        && options.getKinesisPartitionKey() != null

        && options.getKinesisStreamName() != null;

  }



  private static void writeToPubSub(PCollection<KV<byte[], byte[]>> collection) {

    collection

        .apply("Map to PubSub messages", MapElements.via(new MapBytesToPubSubMessage()))

        .apply("Write to PubSub", PubsubIO.writeMessages().to(options.getPubSubTopic()));

  }



  private static void writeToKafka(PCollection<KV<byte[], byte[]>> collection) {

    collection

        .apply("Map to Kafka messages", MapElements.via(new MapKVToString()))

        .apply(

            "Write to Kafka",

            KafkaIO.<Void, String>write()

                .withBootstrapServers(options.getKafkaBootstrapServerAddress())

                .withTopic(options.getKafkaTopic())

                .withValueSerializer(StringSerializer.class)

                .values());

  }



  private static void writeToKinesis(PCollection<KV<byte[], byte[]>> collection) {

    collection

        .apply("Map to byte array for Kinesis", MapElements.via(new MapKVToByteArray()))

        .apply(

            "Write to Kinesis",

            KinesisIO.write()

                .withStreamName(options.getKinesisStreamName())

                .withPartitionKey(options.getKinesisPartitionKey())

                .withAWSClientsProvider(

                    options.getKinesisAwsKey(),

                    options.getKinesisAwsSecret(),

                    Regions.fromName(options.getKinesisAwsRegion())));

  }



  private static class MapKVToString extends SimpleFunction<KV<byte[], byte[]>, String> {

    @Override

    public String apply(KV<byte[], byte[]> input) {

      return String.format(

          "{%s,%s}", Arrays.toString(input.getKey()), Arrays.toString(input.getValue()));

    }

  }



  private static class MapKVToByteArray extends SimpleFunction<KV<byte[], byte[]>, byte[]> {

    @Override

    public byte[] apply(KV<byte[], byte[]> input) {

      return encodeInputElement(input);

    }

  }



  private static class MapBytesToPubSubMessage

      extends SimpleFunction<KV<byte[], byte[]>, PubsubMessage> {

    @Override

    public PubsubMessage apply(KV<byte[], byte[]> input) {

      return new PubsubMessage(encodeInputElement(input), encodeInputElementToMapOfStrings(input));

    }

  }



  private static byte[] encodeInputElement(KV<byte[], byte[]> input) {

    try {

      return CoderUtils.encodeToByteArray(RECORD_CODER, input);

    } catch (CoderException e) {

      throw new RuntimeException(String.format("Couldn't encode element. Exception: %s", e));

    }

  }



  private static Map<String, String> encodeInputElementToMapOfStrings(KV<byte[], byte[]> input) {

    String key = new String(input.getKey(), UTF_8);

    String value = new String(input.getValue(), UTF_8);

    HashMap<String, String> map = new HashMap<>();

    map.put(key, value);

    return map;

  }

}